import mysql.connector
from contextlib import contextmanager

class DatabaseConnection:
    """数据库连接管理类"""
    
    def __init__(self, config):
        self.config = config
        self.connection = None
    
    def connect(self):
        """建立数据库连接"""
        if not self.connection or not self.connection.is_connected():
            self.connection = mysql.connector.connect(**self.config)
        return self.connection
    
    def close(self):
        """关闭数据库连接"""
        if self.connection and self.connection.is_connected():
            self.connection.close()
    
    @contextmanager
    def cursor(self, dictionary=True):
        """创建游标的上下文管理器"""
        cursor = None
        try:
            cursor = self.connect().cursor(dictionary=dictionary)
            yield cursor
        finally:
            if cursor:
                cursor.close()
    
    @contextmanager
    def transaction(self):
        """事务管理的上下文管理器"""
        try:
            yield self.connection
            self.connection.commit()
        except Exception:
            self.connection.rollback()
            raise 